{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RK: this is outdated and may not work.\n",
    "#     Keep this ipynb as a placeholder for zhaires_to_root.py\n",
    "#     Update this after zhaires_to_root.py is complete.\n",
    "\n",
    "from __future__ import annotations\n",
    "\n",
    "from datetime import datetime\n",
    "from logging import getLogger\n",
    "from pathlib import Path\n",
    "import re\n",
    "from typing import Any, Dict, Optional\n",
    "\n",
    "from astropy.coordinates import (\n",
    "    BaseCoordinateFrame,\n",
    "    CartesianRepresentation,\n",
    "    PhysicsSphericalRepresentation,\n",
    ")\n",
    "import astropy.units as u\n",
    "import h5py\n",
    "import numpy\n",
    "\n",
    "# from .generic import CollectionEntry, FieldsCollection, ShowerEvent\n",
    "from ipynb.fs.full.shower_generic import CollectionEntry, FieldsCollection, ShowerEvent\n",
    "\n",
    "# from ..antenna import ElectricField\n",
    "from ipynb.fs.full.antenna import ElectricField\n",
    "from grand.simulation.pdg import ParticleCode\n",
    "from grand.tools.coordinates import ECEF, LTP\n",
    "\n",
    "__all__ = [\"InvalidAntennaName\", \"ZhairesShower\"]\n",
    "\n",
    "\n",
    "logger = getLogger(__name__)\n",
    "\n",
    "\n",
    "class InvalidAntennaName(ValueError):\n",
    "    pass\n",
    "\n",
    "\n",
    "class ZhairesShower(ShowerEvent):\n",
    "    @classmethod\n",
    "    def _check_dir(cls, path: Path) -> bool:\n",
    "        try:\n",
    "            info_file = path.glob(\"*.sry\").__next__()\n",
    "        except StopIteration:\n",
    "            return False\n",
    "        return True\n",
    "\n",
    "    @classmethod\n",
    "    def _from_dir(cls, path: Path) -> ZhairesShower:\n",
    "        if not path.exists():\n",
    "            raise FileNotFoundError(path)\n",
    "\n",
    "        print(\"#### ZhairesShower: load info from dir\")\n",
    "        positions = {}\n",
    "        ant_file = path / \"antpos.dat\"\n",
    "        if ant_file.exists():\n",
    "            pattern = re.compile(\"A([0-9]+)$\")\n",
    "            with ant_file.open() as f:\n",
    "                for line in f:\n",
    "                    if not line:\n",
    "                        continue\n",
    "                    words = line.split()\n",
    "\n",
    "                    match = pattern.search(words[1])\n",
    "                    if match is None:\n",
    "                        raise InvalidAntennaName(words[1])\n",
    "                    antenna = int(match.group(1))\n",
    "\n",
    "                    positions[antenna] = CartesianRepresentation(\n",
    "                        x=float(words[2]) * u.m,\n",
    "                        y=float(words[3]) * u.m,\n",
    "                        z=0 * u.m\n",
    "                        # z = float(words[4]) * u.m\n",
    "                    )\n",
    "                    print(\"### Warning: Forcing antenna height = 0m\")\n",
    "\n",
    "        fields: Optional[FieldsCollection] = None\n",
    "        raw_fields = {}\n",
    "        for field_path in path.glob(\"a*.trace\"):\n",
    "            antenna = int(field_path.name[1:].split(\".\", 1)[0])\n",
    "            logger.debug(f\"Loading trace for antenna {antenna}\")\n",
    "            data = numpy.loadtxt(field_path)\n",
    "            uVm = u.uV / u.m\n",
    "            t = data[:, 0] * u.ns\n",
    "            Ex = data[:, 1] * uVm\n",
    "            Ey = data[:, 2] * uVm\n",
    "            Ez = data[:, 3] * uVm\n",
    "            electric = ElectricField(t, CartesianRepresentation(Ex, Ey, Ez), positions[antenna])\n",
    "            raw_fields[antenna] = CollectionEntry(electric)\n",
    "\n",
    "        if raw_fields:\n",
    "            fields = FieldsCollection()\n",
    "            for key in sorted(raw_fields.keys()):\n",
    "                fields[key] = raw_fields[key]\n",
    "\n",
    "        inp: Dict[str, Any] = {}\n",
    "        print(\"### Setting core\")\n",
    "        inp[\"core\"] = CartesianRepresentation(0, 0, 0, unit=\"m\")\n",
    "        try:\n",
    "            sry_path = path.glob(\"*.sry\").__next__()\n",
    "        except StopIteration:\n",
    "            raise FileNotFoundError(path / \"*.sry\")\n",
    "        else:\n",
    "\n",
    "            def parse_primary(string: str) -> ParticleCode:\n",
    "                return {\"Proton\": ParticleCode.PROTON, \"Iron\": ParticleCode.IRON}[string.strip()]\n",
    "\n",
    "            def parse_quantity(string: str) -> u.Quantity:\n",
    "                words = string.split()\n",
    "                return float(words[0]) * u.Unit(words[1])\n",
    "\n",
    "            def parse_frame_location(string: str) -> BaseCoordinateFrame:\n",
    "                print(\"#### parse_frame_location\")\n",
    "                lat, lon = string.split(\"Long:\")\n",
    "                lat = parse_quantity(lat[:-2])\n",
    "                lon = parse_quantity(lon[:-3])\n",
    "                return ECEF(lat, lon, 0 * u.m, representation_type=\"geodetic\")\n",
    "\n",
    "            def parse_date(string: str) -> datetime:\n",
    "                return datetime.strptime(string.strip(), \"%d/%b/%Y\")\n",
    "\n",
    "            def parse_frame_direction(string: str) -> BaseCoordinateFrame:\n",
    "                print(\"#### parse_frame_direction\")\n",
    "                inp[\"_origin\"] = inp[\"frame\"]\n",
    "\n",
    "                string = string.strip()\n",
    "                if string == \"Local magnetic north\":\n",
    "                    return \"NWU\"\n",
    "                else:\n",
    "                    raise NotImplementedError(string)\n",
    "\n",
    "            def parse_geomagnet_intensity(string: str) -> u.Quantity:\n",
    "                return float(string.split()[0]) << u.uT\n",
    "\n",
    "            def parse_geomagnet_angles(string: str) -> CartesianRepresentation:\n",
    "                intensity = inp[\"geomagnet\"]\n",
    "                inclination, _, _, declination, _ = string.split()\n",
    "                theta = (90 + float(inclination)) << u.deg\n",
    "                inp[\"_declination\"] = float(declination) << u.deg\n",
    "                spherical = PhysicsSphericalRepresentation(theta=theta, phi=0 << u.deg, r=intensity)\n",
    "                return spherical.represent_as(CartesianRepresentation)\n",
    "\n",
    "            def parse_maximum(string: str) -> CartesianRepresentation:\n",
    "                print(\"#### parse_maximum\")\n",
    "                _, _, *xyz = string.split()\n",
    "                x, y, z = map(float, xyz)\n",
    "                return CartesianRepresentation(x * u.km, y * u.km, (z - 2900) * u.km)\n",
    "\n",
    "            converters = (\n",
    "                (\"(Lat\", \"frame\", parse_frame_location),\n",
    "                (\"Date\", \"_obstime\", parse_date),\n",
    "                (\"Primary particle\", \"primary\", parse_primary),\n",
    "                (\"Primary energy\", \"energy\", parse_quantity),\n",
    "                (\"Primary zenith angle\", \"zenith\", parse_quantity),\n",
    "                (\"Primary azimuth angle\", \"azimuth\", parse_quantity),\n",
    "                (\"Zero azimuth direction\", \"frame\", parse_frame_direction),\n",
    "                (\"Geomagnetic field: Intensity:\", \"geomagnet\", parse_geomagnet_intensity),\n",
    "                (\"I:\", \"geomagnet\", parse_geomagnet_angles),\n",
    "                (\"Location of max.(Km)\", \"maximum\", parse_maximum),\n",
    "            )\n",
    "\n",
    "            i = 0\n",
    "            tag, k, convert = converters[i]\n",
    "            with sry_path.open() as f:\n",
    "                for line in f:\n",
    "                    start = line.find(tag)\n",
    "                    if start < 0:\n",
    "                        continue\n",
    "\n",
    "                    inp[k] = convert(line[start + len(tag) + 1 :])\n",
    "                    i = i + 1\n",
    "                    try:\n",
    "                        tag, k, convert = converters[i]\n",
    "                    except IndexError:\n",
    "                        break\n",
    "\n",
    "        origin = inp.pop(\"_origin\")\n",
    "        print(\"# Origin =\", origin)\n",
    "        declination = inp.pop(\"_declination\")\n",
    "        obstime = inp.pop(\"_obstime\")\n",
    "        orientation = inp[\"frame\"]\n",
    "        inp[\"frame\"] = LTP(\n",
    "            location=origin, orientation=orientation, declination=declination, obstime=obstime\n",
    "        )\n",
    "\n",
    "        return cls(fields=fields, **inp)\n",
    "\n",
    "    @classmethod\n",
    "    def _from_datafile(cls, path: Path) -> ZhairesShower:\n",
    "        with h5py.File(path, \"r\") as fd:\n",
    "            if not \"RunInfo.__table_column_meta__\" in fd[\"/\"]:\n",
    "                return super()._from_datafile(path)\n",
    "\n",
    "            for name in fd[\"/\"].keys():\n",
    "                if not name.startswith(\"RunInfo\"):\n",
    "                    break\n",
    "\n",
    "            event = fd[f\"{name}/EventInfo\"]\n",
    "            antennas = fd[f\"{name}/AntennaInfo\"]\n",
    "            traces = fd[f\"{name}/AntennaTraces\"]\n",
    "\n",
    "            fields = FieldsCollection()\n",
    "\n",
    "            pattern = re.compile(\"([0-9]+)$\")\n",
    "            for tag, x, y, z, *_ in antennas:\n",
    "                tag = tag.decode()\n",
    "                antenna = int(pattern.search(tag)[1])\n",
    "                r = CartesianRepresentation(float(x), float(y), float(z), unit=u.m)\n",
    "                tmp = traces[f\"{tag}/efield\"][:]\n",
    "                efield = tmp.view(\"f4\").reshape(tmp.shape + (-1,))\n",
    "                t = numpy.asarray(efield[:, 0], \"f8\") << u.ns\n",
    "                Ex = numpy.asarray(efield[:, 1], \"f8\") << u.uV / u.m\n",
    "                Ey = numpy.asarray(efield[:, 2], \"f8\") << u.uV / u.m\n",
    "                Ez = numpy.asarray(efield[:, 3], \"f8\") << u.uV / u.m\n",
    "                E = (CartesianRepresentation(Ex, Ey, Ez, copy=False),)\n",
    "\n",
    "                fields[antenna] = CollectionEntry(electric=ElectricField(t=t, E=E, r=r))\n",
    "\n",
    "            primary = {\n",
    "                \"Fe^56\": ParticleCode.IRON,\n",
    "                \"Gamma\": ParticleCode.GAMMA,\n",
    "                \"Proton\": ParticleCode.PROTON,\n",
    "            }[event[0, \"Primary\"].decode()]\n",
    "\n",
    "            geomagnet = PhysicsSphericalRepresentation(\n",
    "                theta=float(90 + event[0, \"BFieldIncl\"]) << u.deg,\n",
    "                phi=0 << u.deg,\n",
    "                r=float(event[0, \"BField\"]) << u.uT,\n",
    "            )\n",
    "\n",
    "            try:\n",
    "                latitude = event[0, \"Latitude\"] << u.deg\n",
    "                longitude = event[0, \"Longitude\"] << u.deg\n",
    "                declination = event[0, \"BFieldDecl\"] << u.deg\n",
    "                obstime = datetime.strptime(event[0, \"Date\"].strip(), \"%d/%b/%Y\")\n",
    "            except ValueError:\n",
    "                frame = None\n",
    "            else:\n",
    "                origin = ECEF(latitude, longitude, 0 * u.m, representation_type=\"geodetic\")\n",
    "                frame = LTP(\n",
    "                    location=origin, orientation=\"NWU\", declination=declination, obstime=obstime\n",
    "                )\n",
    "\n",
    "            return cls(\n",
    "                energy=float(event[0, \"Energy\"]) << u.EeV,\n",
    "                zenith=(180 - float(event[0, \"Zenith\"])) << u.deg,\n",
    "                azimuth=-float(event[0, \"Azimuth\"]) << u.deg,\n",
    "                primary=primary,\n",
    "                frame=frame,\n",
    "                core=CartesianRepresentation(0, 0, 2900, unit=\"m\"),\n",
    "                geomagnet=geomagnet.represent_as(CartesianRepresentation),\n",
    "                maximum=CartesianRepresentation(*event[0, \"XmaxPosition\"], unit=\"m\"),\n",
    "                fields=fields,\n",
    "            )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
